#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
音频录制模块
提供录音功能的核心类
"""

import sounddevice as sd
import numpy as np
import wave
import threading
import time
from datetime import datetime
import os

class AudioRecorder:
    def __init__(self, sample_rate=16000, channels=1, dtype=np.float32):
        """
        初始化录音器
        
        Args:
            sample_rate (int): 采样率，默认16000Hz
            channels (int): 声道数，默认1（单声道）
            dtype: 数据类型，默认float32
        """
        self.sample_rate = sample_rate
        self.channels = channels
        self.dtype = dtype
        self.recording = False
        self.audio_data = []
        self.record_thread = None
        
    def start_recording(self):
        """
        开始录音
        """
        if self.recording:
            raise RuntimeError("录音已在进行中")
            
        self.recording = True
        self.audio_data = []
        
        # 启动录音线程
        self.record_thread = threading.Thread(target=self._record_audio)
        self.record_thread.daemon = True
        self.record_thread.start()
        
    def stop_recording(self, output_file=None):
        """
        停止录音并保存文件
        
        Args:
            output_file (str): 输出文件路径，如果为None则自动生成
            
        Returns:
            str: 保存的文件路径
        """
        if not self.recording:
            raise RuntimeError("当前没有在录音")
            
        self.recording = False
        
        # 等待录音线程结束
        if self.record_thread and self.record_thread.is_alive():
            self.record_thread.join(timeout=2.0)
            
        # 生成输出文件名
        if output_file is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            output_file = f"recording_{timestamp}.wav"
            
        # 保存音频文件
        if self.audio_data:
            self._save_audio(output_file)
            return output_file
        else:
            raise RuntimeError("没有录制到音频数据")
            
    def is_recording(self):
        """
        检查是否正在录音
        
        Returns:
            bool: 是否正在录音
        """
        return self.recording
        
    def get_available_devices(self):
        """
        获取可用的音频设备
        
        Returns:
            dict: 设备信息
        """
        try:
            devices = sd.query_devices()
            return devices
        except Exception as e:
            print(f"获取设备信息失败: {e}")
            return None
            
    def _record_audio(self):
        """
        录音线程函数
        """
        try:
            # 使用sounddevice录音
            with sd.InputStream(
                samplerate=self.sample_rate,
                channels=self.channels,
                dtype=self.dtype,
                callback=self._audio_callback
            ) as stream:
                while self.recording:
                    time.sleep(0.1)
                    
        except Exception as e:
            print(f"录音过程中发生错误: {e}")
            self.recording = False
            
    def _audio_callback(self, indata, frames, time, status):
        """
        音频回调函数
        
        Args:
            indata: 输入音频数据
            frames: 帧数
            time: 时间信息
            status: 状态信息
        """
        if status:
            print(f"音频回调状态: {status}")
            
        if self.recording:
            # 将音频数据添加到缓冲区
            self.audio_data.append(indata.copy())
            
    def _save_audio(self, filename):
        """
        保存音频文件
        
        Args:
            filename (str): 文件名
        """
        if not self.audio_data:
            raise ValueError("没有音频数据可保存")
            
        # 合并所有音频数据
        audio_array = np.concatenate(self.audio_data, axis=0)
        
        # 确保目录存在
        os.makedirs(os.path.dirname(filename) if os.path.dirname(filename) else '.', exist_ok=True)
        
        # 保存为WAV文件
        with wave.open(filename, 'wb') as wf:
            wf.setnchannels(self.channels)
            wf.setsampwidth(4)  # float32 = 4 bytes
            wf.setframerate(self.sample_rate)
            
            # 转换为16位整数格式保存
            audio_int16 = (audio_array * 32767).astype(np.int16)
            wf.setsampwidth(2)  # 16位 = 2 bytes
            wf.writeframes(audio_int16.tobytes())
            
        print(f"音频已保存到: {filename}")

# 测试函数
def test_recorder():
    """
    测试录音功能
    """
    recorder = AudioRecorder()
    
    print("开始录音...")
    recorder.start_recording()
    
    # 录音5秒
    time.sleep(5)
    
    print("停止录音...")
    output_file = recorder.stop_recording()
    print(f"录音已保存到: {output_file}")

if __name__ == "__main__":
    test_recorder()